import torch
import matplotlib.pyplot as plt
from deoldify.device_id import DeviceId
from deoldify import device
from deoldify.visualize import *
import warnings
from PIL import Image
import tempfile
import gradio as gr
import numpy as np


# 忽略警告信息
warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")

# 设置使用GPU进行加速计算，如果可用
device.set(device=DeviceId.GPU0 if torch.cuda.is_available() else DeviceId.CPU)

# 设置matplotlib绘图风格为深色背景
plt.style.use('dark_background')

# 设置PyTorch使用benchmark模式，以加速GPU运算
torch.backends.cudnn.benchmark = True

# 创建一个图像上色器实例
colorizer = get_image_colorizer(artistic=True)

# 定义一个函数，用于将黑白图像上色
def colorize_image(image, render_factor):

    # 检查是否提供了图像
    if image is None:
        raise ValueError("No image provided")

    # 将PIL Image转换为numpy数组
    image = Image.fromarray(image)
    
    # 将图像保存为临时文件
    with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as f:
        image_path = f.name
    image.save(image_path, "JPEG")
    
    # 使用滑动条的值来进行上色并返回结果路径
    result_path = colorizer.plot_transformed_image(path=image_path, render_factor=render_factor, compare=True)
    
    # 读取上色后的图像
    colorized_image = Image.open(result_path)
    
    # 将PIL Image转换为numpy数组并返回
    colorized_image_np = np.array(colorized_image)
    return colorized_image_np

# 创建Gradio界面
def gradio_interface():
  iface = gr.Interface(
     fn=colorize_image,
     inputs=[
        gr.Image(label="Upload Black and White Image"),
        gr.Slider(minimum=0, maximum=50, step=1, value=0, label="Render Factor")
    ],
     outputs=gr.Image(label="Colorized Image"),
     title="DeOldify lmage Colorizer",
     description="上传一张黑白图片并通过滑动条来进行上色。"
)

# 运行应用
iface.launch(share=True)
